import ctypes
import numpy as np

# 加载RGA库
lib = ctypes.cdll.LoadLibrary("/data/detect/lib/libmyrga.so")

WIDTH=2560
HEIGHT=1440

# prototypes
lib.nv12_to_rgb_init.argtypes = [ctypes.c_int, ctypes.c_int]
lib.nv12_to_rgb_init.restype = ctypes.c_int

lib.nv12_alloc_src.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_void_p)]
lib.nv12_alloc_src.restype = ctypes.c_int  # returns allocated size

lib.nv12_free_src.argtypes = [ctypes.c_int, ctypes.c_size_t, ctypes.c_void_p]
lib.nv12_free_src.restype = None

lib.nv12_to_rgb_convert.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_int]
lib.nv12_to_rgb_convert.restype = ctypes.c_int

lib.nv12_get_dst_fd.restype = ctypes.c_int
lib.nv12_get_dst_virt.restype = ctypes.c_void_p
lib.nv12_get_dst_size.restype = ctypes.c_size_t

lib.nv12_to_rgb_deinit.restype = None

# 初始化（例如 3840x2160 支持 4K）
if lib.nv12_to_rgb_init(WIDTH, HEIGHT) != 0:
    raise RuntimeError("init failed")

# 分配 src NV12
w, h = WIDTH, HEIGHT
src_fd = ctypes.c_int()
src_virt = ctypes.c_void_p()
allocated_size = lib.nv12_alloc_src(w, h, ctypes.byref(src_fd), ctypes.byref(src_virt))
if allocated_size <= 0:
    raise RuntimeError("alloc src failed")
print("src_fd:", src_fd.value, "src_virt:", hex(src_virt.value), "size:", allocated_size)

# 举例：把 input.nv12 写入 src_virt
with open("output_nv12.yuv", "rb") as f:
    data = f.read()
    # 注意：确保 data 长度 <= allocated_size
    buf_type = ctypes.c_char * allocated_size
    buf = buf_type.from_address(src_virt.value)
    buf[:len(data)] = data

# 调用转换
ret = lib.nv12_to_rgb_convert(src_fd.value, w, h)
if ret != 0:
    raise RuntimeError("convert failed")

# 读取 dst
dst_fd = lib.nv12_get_dst_fd()
dst_ptr = lib.nv12_get_dst_virt()
dst_size = lib.nv12_get_dst_size()
print("dst_fd:", dst_fd, "dst_virt:", hex(dst_ptr), "dst_size:", dst_size)

# 将 RGB 数据保存 (注意：RGB 不含头，仅 raw rgb24)
dst_buf_type = ctypes.c_ubyte * dst_size
dst_buf = dst_buf_type.from_address(dst_ptr)
with open("output.rgb", "wb") as f:
    f.write(bytes(dst_buf))
    
# 清理
lib.nv12_free_src(src_fd.value, allocated_size, src_virt.value)
lib.nv12_to_rgb_deinit()